import numpy as np
from torchvision import datasets, transforms


def mnist_iid(dataset, num_users):
    """Sample IID client data from MNIST dataset."""
    num_items = int(len(dataset) / num_users)
    dict_users, all_idxs = {}, [i for i in range(len(dataset))]
    for i in range(num_users):
        dict_users[i] = set(np.random.choice(all_idxs, num_items, replace=False))
        all_idxs = list(set(all_idxs) - dict_users[i])
    return dict_users


def mnist_noniid(dataset, num_users):
    """Sample non-IID client data from MNIST dataset."""
    num_shards, num_imgs = 200, 300
    idx_shard = [i for i in range(num_shards)]
    dict_users = {i: np.array([], dtype='int64') for i in range(num_users)}
    idxs = np.arange(num_shards * num_imgs)
    labels = dataset.targets.numpy()

    idxs_labels = np.vstack((idxs, labels))
    idxs_labels = idxs_labels[:, idxs_labels[1, :].argsort()]
    idxs = idxs_labels[0, :]

    for i in range(num_users):
        rand_set = set(np.random.choice(idx_shard, 2, replace=False))
        idx_shard = list(set(idx_shard) - rand_set)
        for rand in rand_set:
            dict_users[i] = np.concatenate(
                (dict_users[i], idxs[rand * num_imgs:(rand + 1) * num_imgs]),
                axis=0
            )
    return dict_users


def cifar_iid(dataset, num_users):
    """Sample IID client data from CIFAR10 dataset."""
    return mnist_iid(dataset, num_users)


def cifar_noniid(dataset, num_users):
    """Sample non-IID client data from CIFAR10 dataset."""
    num_shards, num_imgs = 200, 250
    idx_shard = [i for i in range(num_shards)]
    dict_users = {i: np.array([], dtype='int64') for i in range(num_users)}
    idxs = np.arange(num_shards * num_imgs)
    labels = np.array(dataset.targets)

    idxs_labels = np.vstack((idxs, labels))
    idxs_labels = idxs_labels[:, idxs_labels[1, :].argsort()]
    idxs = idxs_labels[0, :]

    for i in range(num_users):
        rand_set = set(np.random.choice(idx_shard, 2, replace=False))
        idx_shard = list(set(idx_shard) - rand_set)
        for rand in rand_set:
            dict_users[i] = np.concatenate(
                (dict_users[i], idxs[rand * num_imgs:(rand + 1) * num_imgs]),
                axis=0
            )
    return dict_users


def fmnist_iid(dataset, num_users):
    """Sample IID client data from FashionMNIST dataset."""
    return mnist_iid(dataset, num_users)


def fmnist_noniid(dataset, num_users):
    """Sample non-IID client data from FashionMNIST dataset."""
    return mnist_noniid(dataset, num_users)


def svhn_iid(dataset, num_users):
    """Sample IID client data from SVHN dataset."""
    return cifar_iid(dataset, num_users)


def svhn_noniid(dataset, num_users):
    """Sample non-IID client data from SVHN dataset."""
    num_imgs = 300
    num_shards = len(dataset) // num_imgs
    shards_per_user = 2

    if num_shards < num_users * shards_per_user:
        raise ValueError(
            f"Insufficient shards: need {num_users * shards_per_user}, have {num_shards}"
        )

    idx_shard = list(range(num_shards))
    dict_users = {i: np.array([], dtype='int64') for i in range(num_users)}
    labels = np.array(dataset.labels).squeeze()
    idxs = np.arange(len(dataset))

    idxs_labels = np.vstack((idxs, labels))
    sorted_indices = idxs_labels[:, idxs_labels[1, :].argsort()]
    sorted_idxs = sorted_indices[0, :]

    np.random.shuffle(idx_shard)
    for i in range(num_users):
        selected_shards = idx_shard[i * shards_per_user: (i + 1) * shards_per_user]
        for shard in selected_shards:
            start = shard * num_imgs
            end = (shard + 1) * num_imgs
            dict_users[i] = np.concatenate(
                (dict_users[i], sorted_idxs[start:end]),
                axis=0
            )
    return dict_users

